# -*- coding: utf-8 -*-
# @File     : train_cfg.py
# @Author   : bingjia
# @Time     : 2020/9/22 9:30
# @Desc     : 训练配置文件
import os
from easydict import EasyDict as edict

ROOT_DIR = os.path.abspath(os.path.dirname(os.path.dirname(__file__)))

__C = edict()
train_cfg = __C

"""
训练设置
"""
__C.hyp = ''

__C.bucket = ''

"result.txt文件名字"
__C.name = ''

"训练设备，cpu或者gpu 0 1 2"
__C.device = ''

"权重文件"
__C.weights = ''

"迭代次数"
__C.epochs = 300

__C.world_size = 1

"批处理大小"
__C.batch_size = 12

"DDP参数"
__C.local_rank = -1

"使用矩形训练"
__C.rect = False

"使用超参数"
__C.evolve = False

"只保存最终的checkpoint文件"
__C.nosave = False

"只进行最有一个epoch的测试"
__C.notest = False

"使用恢复训练"
__C.recover = True

" 使用同步批处理"
__C.sync_bn = False

"使用多尺度训练"
__C.multi_scale = True

"使训练为单类数据集"
__C.single_cls = False

"使用缓存图像文件快速训练"
__C.cache_images = False

"禁用autoanchor检查"
__C.noautoanchor = False

"batch size大小"
__C.total_batch_size = 8

"训练、测试图片尺度"
__C.img_size = [640, 640]

"训练数据配置"
__C.data = os.path.join(ROOT_DIR, 'data', 'pv.yaml')

"选用训练模型"
__C.cfg = os.path.join(ROOT_DIR, 'models', 'yolov5s.yaml')

"模型生成路径"
__C.model_save_path = os.path.join(ROOT_DIR, 'runs')


"超参数"
hyp = {'optimizer': 'SGD',  # ['adam', 'SGD', None] if none, default is SGD
       'lr0': 0.01,  # initial learning rate (SGD=1E-2, Adam=1E-3)
       'momentum': 0.937,  # SGD momentum/Adam beta1
       'weight_decay': 5e-4,  # optimizer weight decay
       'giou': 0.05,  # giou loss gain
       'cls': 0.5,  # cls loss gain
       'cls_pw': 1.0,  # cls BCELoss positive_weight
       'obj': 1.0,  # obj loss gain (*=img_size/320 if img_size != 320)
       'obj_pw': 1.0,  # obj BCELoss positive_weight
       'iou_t': 0.20,  # iou training threshold
       'anchor_t': 4.0,  # anchor-multiple threshold
       'fl_gamma': 0.0,  # focal loss gamma (efficientDet default is gamma=1.5)
       'hsv_h': 0.015,  # image HSV-Hue augmentation (fraction)
       'hsv_s': 0.7,  # image HSV-Saturation augmentation (fraction)
       'hsv_v': 0.4,  # image HSV-Value augmentation (fraction)
       'degrees': 0.0,  # image rotation (+/- deg)
       'translate': 0.0,  # image translation (+/- fraction)
       'scale': 0.5,  # image scale (+/- gain)
       'shear': 0.0}  # image shear (+/- deg)
